import pandas as pd
import os
import numpy as np
import info
from StrategyModule import Strategy
import datetime
from TradeRecordV1 import TradeRecord
# from TradeRecordModule import TradeRecord
import matplotlib.pyplot as plt
from GetExternalDataSourceModule_CLIENT import GetExternalDataSource_CLIENT

class Environment:
    def __init__(self, data_name: str, data_address: str, data_location:str, trade_price_name:str, daily_clearing_price_name:str, max_lookback_bars: int, max_trade_times: int, trade_unit: int, stop_win: float, stop_loss: float,
                 x_min_price: int, long_freeze_period: int, short_freeze_period:int):

        # 初始化服务器连接模块
        self.connector = GetExternalDataSource_CLIENT()
        # 种类名称
        self.product_name = data_name.split('_')[0]
        # 频率
        self.freq = data_name.split('_')[1]
        # 交易价格字段名称
        self.trade_price_name = trade_price_name
        # 清算价格字段名称
        self.daily_clearing_price_name = daily_clearing_price_name
        # self.ZL_path = info.DATAPATH.dataRootPath + '\\' + 'future_trade_data' + '\\' + data_address
        # self.ZL_info = pd.read_pickle(info.DATAPATH.dataRootPath + '\\' + 'future_basic_data' + '\\' + 'ZL.pkl')[self.product_name].dropna()
        # 载入主力合约信息
        self.ZL_info = self.connector.get_file('ZL', info.DATAPATH.server_root_path + '/' + 'future_basic_data' + '/' + 'ZL.pkl')[self.product_name].dropna()
        # 获得更换合约日
        self.rotate_contract_signals = self.ZL_info !=self.ZL_info.shift(1)
        # 载入字典
        self.calendar = self.connector.get_file('tradingCalendar', info.DATAPATH.server_root_path + '/' +
                                                'calendar' + '/' + 'tradingCalendar.pkl')
        self.calendar.index = self.calendar.TradingDate.tolist()

        if data_location=='local':
            # 从本地载入主力数据
            self.raw_trade_data = pd.read_pickle(
                info.DATAPATH.dataRootPath + '\\' + 'future_trade_data' + '\\' +  data_address + '\\' +
                self.product_name + '_' + self.freq + '.pkl')
        elif data_location=='remote':
            # 从服务器载入主力数据
            self.raw_trade_data = self.connector.get_file(self.product_name + '_' + self.freq,
                                                          remote_path=info.DATAPATH.server_root_path + '/' +
                                                                      'future_trade_data' + '/' +  data_address + '/' +
                                                                      self.product_name + '_' + self.freq + '.pkl')

            self.raw_trade_data = self.raw_trade_data[
                ~self.raw_trade_data.index.duplicated(keep='first')]
            self.raw_trade_data = self.raw_trade_data.dropna()

        if self.freq != 'D':
            if not 'date' in self.raw_trade_data.columns.tolist():
                self.raw_trade_data['date'] = [str(x)[:10] for x in self.raw_trade_data.index.tolist()]

                self.raw_trade_data = self.raw_trade_data.loc[
                                        np.in1d(self.raw_trade_data.date.tolist(), self.calendar.index.tolist()), :]
                self.raw_trade_data['date'] = [
                    str(x)[:10] if '16:00:00' >= str(x.time()) >= '00:00:00' else self.calendar.shift(-1).loc[
                        str(x)[:10], 'TradingDate'] for x in self.raw_trade_data.index.tolist()]
        else:
            self.raw_trade_data['date'] = [str(x)[:10] for x in self.raw_trade_data.index.tolist()]

        # 载入复权因子
        self.price_rate = self.connector.get_file('PriceRate', info.DATAPATH.server_root_path +
                                                  '/' + 'future_basic_data' + '/' + 'PriceRate.pkl')[self.product_name]
        # 载入最大回看bar数
        self.max_lookback_bars = max_lookback_bars
        # 载入交易单元
        self.trade_unit = trade_unit
        # 止盈
        self.stop_win = stop_win
        # 止损
        self.stop_loss = stop_loss
        # 几个最小价格变动
        self.x_min_price = x_min_price
        # 每日最多交易次数
        self.max_trade_times = max_trade_times
        # 买入冷冻期
        self.long_freeze_period = long_freeze_period
        # 卖出冷冻期
        self.short_freeze_period = short_freeze_period
        print('Enviroment initiates.')

    def backtest_setting(self, start_date: str, end_date: str, clearing_time: list):
        # 每日回测开始
        if not 'D' in self.freq:
            self._backtest_start_date = datetime.datetime.strptime(start_date + ' 09:00:00', '%Y-%m-%d %H:%M:%S')
        else:
            self._backtest_start_date = datetime.datetime.strptime(start_date + ' 00:00:00', '%Y-%m-%d %H:%M:%S')
        # 每日回测结束
        self._backtest_end_date = datetime.datetime.strptime(end_date + ' 16:00:00', '%Y-%m-%d %H:%M:%S')
        # 清算时间段
        self._clearing_time = clearing_time
        #
        if len(clearing_time)!=0:
            self._clearing_daily = True
        else:
            self._clearing_daily=False


    def setExternalData(self, external_data):
        # 外部数据
        self.external_data = external_data

    def getExternalData(self, ):
        assert self.external_data
        return self.external_data

    def backtest_running(self):
        # 主力合约原始行情
        ZL = self.raw_trade_data[self.raw_trade_data.ZL]
        # 换仓日次主力合约原始行情
        ZL_last = self.raw_trade_data[~self.raw_trade_data.ZL]
        self.trade_recorder = TradeRecord(
                                            trade_unit=self.trade_unit,
                                            stop_win=self.stop_win,
                                            stop_loss=self.stop_loss,
                                            raw_data_ZL=ZL,
                                            raw_data_ZL_last=ZL_last,
                                            x_min_price=self.x_min_price,
                                            rotate_signal=self.rotate_contract_signals,
                                            clearing_time=self._clearing_time,
                                            rotate_trade_name=self.trade_price_name,
                                            max_trade_times=self.max_trade_times,
                                            long_freeze_period = self.long_freeze_period,
                                            short_freeze_period = self.short_freeze_period,
                                            freq = self.freq
                                          )
        # 历史bar
        previous_bars = ZL.loc[:self._backtest_start_date, :].iloc[-self.max_lookback_bars-1:-1, :]
        # 复权历史bar
        total_bars_adjust = self.getAdjustPrice(ZL)
        backtesting_bars = ZL.loc[self._backtest_start_date:self._backtest_end_date, :]
        total_bars = pd.concat([previous_bars, backtesting_bars, ])
        total_bars_adjust = total_bars_adjust.loc[total_bars.index.tolist(), :]

        for bar_datetime in backtesting_bars.index.tolist():
            current_bar_datetime = bar_datetime
            current_bar_contract_name = backtesting_bars.loc[bar_datetime, 'StockID']
            current_bar_trade_price = backtesting_bars.loc[bar_datetime, self.trade_price_name]
            current_bar_clearing_price = backtesting_bars.loc[bar_datetime, self.daily_clearing_price_name]

            current_backtest_block = total_bars_adjust.loc[:current_bar_datetime, :].iloc[-self.max_lookback_bars-1:-1, :]
            current_bar_last_price = backtesting_bars.loc[bar_datetime, 'yclose']
            current_bar_trade_date = backtesting_bars.loc[bar_datetime, 'date']

            self.strategy.strategy_calculate(current_backtest_block)
            self.strategy.trade_record(
                tradeRecord=self.trade_recorder,
                current_bar_contract_name=current_bar_contract_name,
                close_position_price=current_bar_trade_price,
                open_position_price=current_bar_trade_price,
                daily_clearing_position_price=current_bar_clearing_price,
                last_bar_price=current_bar_last_price,
                current_bar_datetime=current_bar_datetime,
                clearing_daily=self._clearing_daily,
                freq=self.freq,
                current_bar_trade_date=current_bar_trade_date
            )


    def cut_data_to_single_day(self, df, keyword):
        dates = df.loc[:, keyword].unique().tolist()
        daily_data_dict = {}
        for date in dates:
            daily_data_dict[date] = df.loc[df.loc[:, keyword]==date, :]
        return daily_data_dict

    def getAdjustPrice(self, raw_trade_data):
        # price_rate = self.price_rate.loc[raw_trade_data.index.tolist()]
        raw_trade_data_adjust = raw_trade_data.copy(deep=True)
        price_rate = self.price_rate.to_frame()
        price_rate['date'] = [str(x)[:10] for x in price_rate.index.tolist()]
        raw_trade_data_adjust=pd.merge(raw_trade_data_adjust, price_rate, how='left', on='date')
        raw_trade_data_adjust[self.product_name] = raw_trade_data_adjust[self.product_name] - raw_trade_data_adjust[self.product_name][0]
        raw_trade_data_adjust['close'] = raw_trade_data_adjust['close'] + raw_trade_data_adjust[self.product_name]
        raw_trade_data_adjust['open'] = raw_trade_data_adjust['open'] + raw_trade_data_adjust[self.product_name]
        raw_trade_data_adjust['high'] = raw_trade_data_adjust['high'] + raw_trade_data_adjust[self.product_name]
        raw_trade_data_adjust['low'] = raw_trade_data_adjust['low'] + raw_trade_data_adjust[self.product_name]
        raw_trade_data_adjust['yclose'] = raw_trade_data_adjust['close'].shift(1).fillna(0)
        raw_trade_data_adjust['yclose'] = raw_trade_data_adjust['yclose'].fillna(0)
        raw_trade_data_adjust.index = raw_trade_data.index.tolist()
        return raw_trade_data_adjust


    def add_strategy(self, strategy: Strategy):
        self.strategy = strategy

